from core.function import Function


class LeakyReLU(Function):
    def __init__(self, slope):
        """
        具有指定负值斜率的泄漏修正线性单元（Leaky ReLU）函数。

        Args:
            slope (float): 负值的斜率。
        """
        self.slope = slope

    def forward(self, x):
        """
        Leaky ReLU 的前向传播。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: 应用 Leaky ReLU 后的输出变量。
        """
        y = x.copy()
        y[x <= 0] *= self.slope
        return y

    def backward(self, gy):
        """
        Leaky ReLU 的反向传播。

        Args:
            gy (Variable): 输出的梯度。

        Returns:
            Variable: 输入的梯度。
        """
        x, = self.inputs
        mask = (x.data > 0).astype(gy.dtype)
        mask[mask <= 0] = self.slope
        gx = gy * mask
        return gx

    def leaky_relu(self, x):
        """
        对输入变量应用 Leaky ReLU。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: 应用 Leaky ReLU 后的输出变量。
        """
        return self(x)
